. path.sh
export CUDA_DEVICE_MAX_CONNECTIONS=1
export NCCL_LAUNCH_TIMEOUT=30
export NCCL_SOCKET_IFNAME=eth1
#export NCCL_IB_GID_INDEX=3
#export NCCL_IB_HCA=mlx5_2:1
#export NCCL_IB_SL=3
#export NCCL_CHECK_DISABLE=1
export NCCL_P2P_DISABLE=0
export NCCL_IB_DISABLE=1
#export NCCL_DEBUG=INFO
export NCCL_LL_THRESHOLD=16384
#export NCCL_IB_CUDA_SUPPORT=1

# export NVTE_APPLY_QK_LAYER_SCALING=1

TOKENIZER_TYPE=GPT2BPETokenizer

# HF_LLAMA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/Qwen/Qwen2-Audio-7B-Instruct
HF_LLAMA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/Qwen/Qwen2.5-7B

merge_path=$HF_LLAMA_PATH/merges.txt
vocab_path=$HF_LLAMA_PATH/vocab.json


GPUS_PER_NODE=4
# Change for mltinode config
MASTER_ADDR=localhost
MASTER_PORT=12324
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))


TP=2
PP=1
CP=1
MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=2
SEQ_LENGTH=1024
# SAVE_PREFIX=moeconformer_llm_poc
SAVE_PREFIX=debug


HIDDEN_SIZE=3584
FFN_HIDDEN_SIZE=18944
NUM_LAYERS=28
NUM_HEADS=28
VOCAB_SIZE=152064


HIDDEN_SIZE=384
FFN_HIDDEN_SIZE=1024
NUM_LAYERS=8
NUM_HEADS=12

VOCAB_SIZE=152064

#DATA_PATH=/apdcephfs_qy3/share_976139/users/joyounglv/audiollama/data/aishell1/valid_asr_aishell1.json
#DATA_PATH=/apdcephfs_qy3/share_976139/users/joyounglv/audiollama/data/aishell1/train_asr_aishell1.json
#VALID_DATA_PATH=/apdcephfs_qy3/share_976139/users/joyounglv/audiollama/data/aishell1/valid_asr_aishell1.json
#DATA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/audio_llm/raw_data/processed/train_sub_asr_zhen_train_20240322_cn_slides_1100h.jsonl
# VALID_DATA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/audio_llm/raw_data/processed/valid_sub_asr_zhen_train_20240322_cn_slides_1100h.jsonl
#VALID_DATA_PATH=audio_data/valid_20240322_cn_slides_1100h_asr_multi.jsonl
#DATA_PATH=audio_data/train_sub_asr_zhen_train_20240322_cn_slides_asr_multi.jsonl
# DATA_PATH="audio_data/train_ver250425/train_20250425_llm_asr_3kh_remove_repeat.jsonl,audio_data/train_20231206_slidespeech_487h.jsonl"
# VALID_DATA_PATH="audio_data/valid_20240322_cn_slides_1100h.jsonl,audio_data/test_20241119_slidespeech.jsonl"
DATA_PATH=debug.input.jsonl
VALID_DATA_PATH=debug.input.jsonl

#DATA_PATH=$VALID_DATA_PATH
#VALID_DATA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/audio_llm/raw_data/processed/test_asr_zhen_test_20220221_ailab_5h_org_fbankhires.jsonl
#CHECKPOINT_PATH=qwen2-audio-7b-megatron-TP${TP}-PP${PP}-TE
# CHECKPOINT_PATH=qwen2-audio-7b-instruct-trained-megatron-TP${TP}-PP${PP}-TE
CHECKPOINT_PATH=qwen2.5-7b-megatron-for-audio-TP${TP}-PP${PP}-TE

. utils/parse_options.sh

SAVE_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/exp/exp_audiollm/${SAVE_PREFIX}_TP${TP}_PP${PP}_CP${CP}_MBZ${MICRO_BATCH_SIZE}_GBSZ${GLOBAL_BATCH_SIZE}_seq${SEQ_LENGTH}

#SAVE_PATH=exp_audiollm/moeconformerllm_asr_llmasr_zhen_slides_specaug_250425_TP2_PP1_CP1_MBZ1_GBSZ128_seq1024


tensorboard_output=$SAVE_PATH/tensorboard

mkdir -p $SAVE_PATH/tensorboard
# 判断save_path是否有iter_0000001，没有的话则从checkpoint复制过去

# if [ ! -d "$SAVE_PATH/iter_0000001" ] && [ "$NODE_RANK" -eq 0 ]; then
#     echo "$SAVE_PATH/iter_0000001 not found, copying from $CHECKPOINT_PATH"
#     cp -r $CHECKPOINT_PATH/* $SAVE_PATH/
# fi

cp $0 $SAVE_PATH/

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"
      
MODEL_ARGS="
    --num-layers $NUM_LAYERS \
    --hidden-size $HIDDEN_SIZE \
    --ffn-hidden-size $FFN_HIDDEN_SIZE \
    --num-attention-heads $NUM_HEADS \
    --group-query-attention \
    --num-query-groups 4 \
    --norm-epsilon 1e-6 \
    --seq-length $SEQ_LENGTH \
    --rotary-base 1000000 \
    --max-position-embeddings $SEQ_LENGTH \
    --attention-dropout 0 \
    --hidden-dropout 0 \
    --use-rotary-position-embeddings \
    --normalization RMSNorm \
    --no-position-embedding \
    --disable-bias-linear \
    --swiglu \
    --untie-embeddings-and-output-weights \
    --make-vocab-size-divisible-by $((152064 / TP)) \
    --add-qkv-bias
"

OUTPUT_ARGS="
    --log-interval 1 \
    --save-interval 1 \
    --eval-interval 100 \
    --log-throughput \
    --log-memory-to-tensorboard \
    --log-timers-to-tensorboard \
    --tensorboard-dir $tensorboard_output \
    --eval-iters 6
"

    #--make-vocab-size-divisible-by 156032 \

# torchrun $DISTRIBUTED_ARGS megatron/core/models/multimodal/moe_conformer_llm.py \
torchrun $DISTRIBUTED_ARGS examples/audiollm/finetune_moellm.py \
    $MODEL_ARGS \
    $OUTPUT_ARGS \
    --micro-batch-size $MICRO_BATCH_SIZE \
    --global-batch-size $GLOBAL_BATCH_SIZE \
    --tensor-model-parallel-size $TP \
    --pipeline-model-parallel-size $PP \
    --context-parallel-size $CP \
    --extra-equivalent-layers 16 \
    --lr-warmup-iters 0 \
    --lr 1e-5 \
    --min-lr 1e-6 \
    --lr-decay-iters 9000 \
    --lr-decay-style cosine \
    --weight-decay 0.01 \
    --bf16 \
    --ckpt-format torch \
    --train-iters 100000 \
    --tokenizer-type $TOKENIZER_TYPE \
    --vocab-file $vocab_path \
    --merge-file $merge_path \
    --recompute-activations \
    --recompute-granularity selective \
    --use-distributed-optimizer \
    --optimizer hybridadam \
    --save $SAVE_PATH \
    --load $CHECKPOINT_PATH \
    --dataloader-type external \
    --num-workers 4 \
    --no-load-optim \
    --no-load-rng \
    --train-data-path $DATA_PATH \
    --valid-data-path $VALID_DATA_PATH 

    #--sequence-parallel \

        # --train-data-path /apdcephfs_qy3/share_976139/users/joyounglv/audiollama/data/aishell1/train_asr_aishell1.json \
    # --load ./tmp_whisper_noinit
#    --load ./tmp_whisper
    # --load whisper-large-v3-megatron-TP${TP}-TE


